import cv2
import numpy as np
from onnxruntime import InferenceSession

class YOLOv8Detector:
    def __init__(self, onnx_path, class_names, conf_threshold=0.5, iou_threshold=0.5):
        """
        初始化YOLOv8 ONNX检测器
        :param onnx_path: ONNX模型路径
        :param conf_threshold: 置信度阈值
        :param iou_threshold: NMS的IoU阈值
        """
        self.onnx_path = onnx_path
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        
        # 加载模型
        self.session = InferenceSession(self.onnx_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [output.name for output in self.session.get_outputs()]
        
        # 类别名称
        self.class_names = class_names
    
    def preprocess_image(self, image, target_size=(640, 640)):
        """预处理图像（缩放+填充）"""
        h, w = image.shape[:2]
        scale = min(target_size[0] / w, target_size[1] / h)
        new_w, new_h = int(w * scale), int(h * scale)
        resized = cv2.resize(image, (new_w, new_h))
        
        # 填充黑边
        padded = np.full((target_size[1], target_size[0], 3), 114, dtype=np.uint8)
        pad_top = (target_size[1] - new_h) // 2
        pad_left = (target_size[0] - new_w) // 2
        padded[pad_top:pad_top+new_h, pad_left:pad_left+new_w] = resized
        
        # 转换为模型需要的格式
        padded = padded.transpose(2, 0, 1)  # CHW
        padded = padded[::-1]  # BGR to RGB
        padded = padded / 255.0  # 归一化
        padded = np.ascontiguousarray(padded, dtype=np.float32)
        
        return padded[np.newaxis, ...], (pad_left, pad_top, scale)
    
    @staticmethod
    def nms(boxes, scores, iou_threshold=0.5):
        """非极大值抑制"""
        order = scores.argsort()[::-1]
        keep = []
        
        while order.size > 0:
            i = order[0]
            keep.append(i)
            
            # 计算当前框与其他框的IoU
            xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
            yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
            xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
            yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
            
            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            intersection = w * h
            
            area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
            area_j = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
            union = area_i + area_j - intersection
            
            iou = intersection / union
            
            # 保留IoU小于阈值的框
            inds = np.where(iou <= iou_threshold)[0]
            order = order[inds + 1]
        
        return keep
    
    def detect(self, image):
        """
        执行目标检测
        :param image: 输入图像 (OpenCV格式)
        :return: 检测结果列表，每个元素为[x1, y1, x2, y2, conf, cls_id]
        """
        # 预处理
        preprocessed, (pad_left, pad_top, scale) = self.preprocess_image(image)
        
        # 推理
        outputs = self.session.run(self.output_names, {self.input_name: preprocessed})
        
        # 处理不同输出格式
        if len(outputs) == 1:
            # YOLOv8官方输出格式 [1,84,8400]
            predictions = np.squeeze(outputs[0]).T  # [8400,84]
            
            # 提取置信度和类别
            scores = np.max(predictions[:, 4:], axis=1)
            class_ids = np.argmax(predictions[:, 4:], axis=1)
            
            # 提取并转换边界框 (cx,cy,w,h) -> (x1,y1,x2,y2)
            boxes = predictions[:, :4]
            boxes[:, 0] = boxes[:, 0] - boxes[:, 2] / 2  # x1
            boxes[:, 1] = boxes[:, 1] - boxes[:, 3] / 2  # y1
            boxes[:, 2] = boxes[:, 0] + boxes[:, 2]      # x2
            boxes[:, 3] = boxes[:, 1] + boxes[:, 3]      # y2
            
            # 过滤低置信度
            mask = scores > self.conf_threshold
            boxes = boxes[mask]
            scores = scores[mask]
            class_ids = class_ids[mask]
            
            # 应用NMS
            if len(boxes) > 0:
                keep = self.nms(boxes, scores, self.iou_threshold)
                detections = np.concatenate([
                    boxes[keep],
                    scores[keep, None],
                    class_ids[keep, None]
                ], axis=1)
            else:
                detections = np.empty((0, 6))
        else:
            boxes = outputs[0][0]  # (num_detections, 4)
            scores = outputs[1][0]  # (num_detections,)
            class_ids = outputs[2][0]  # (num_detections,)
            
            # 过滤低置信度
            mask = scores > self.conf_threshold
            boxes = boxes[mask]
            scores = scores[mask]
            class_ids = class_ids[mask]
            
            # 应用NMS
            if len(boxes) > 0:
                keep = self.nms(boxes, scores, self.iou_threshold)
                detections = np.concatenate([
                    boxes[keep],
                    scores[keep, None],
                    class_ids[keep, None]
                ], axis=1)
            else:
                detections = np.empty((0, 6))
        
        # 转换坐标到原始图像空间
        results = []
        for detection in detections:
            if len(detection) == 6:
                x1, y1, x2, y2, conf, cls_id = detection
            else:
                x1, y1, x2, y2 = detection[:4]
                conf = detection[4]
                cls_id = detection[5] if len(detection) > 5 else 0
            
            # 转换坐标
            x1 = int((x1 - pad_left) / scale)
            y1 = int((y1 - pad_top) / scale)
            x2 = int((x2 - pad_left) / scale)
            y2 = int((y2 - pad_top) / scale)
            
            # 确保坐标有效
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
            
            results.append([x1, y1, x2, y2, conf, int(cls_id)])
        
        return results
    
    def visualize(self, image, detections):
        """可视化检测结果"""
        for detection in detections:
            x1, y1, x2, y2, conf, cls_id = detection
            
            # 安全获取类别名称
            cls_name = self.class_names[cls_id] if cls_id < len(self.class_names) else f"unknown_{cls_id}"
            
            label = f"{cls_name} {conf:.2f}"
            
            # 绘制
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.rectangle(image, (x1, y1 - 20), (x1 + len(label) * 10, y1), (0, 255, 0), -1)
            cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
        
        return image

def main():
    # 初始化检测器
    detector = YOLOv8Detector(
        onnx_path="E:/position.onnx",  # 替换为你的模型路径
        class_names=['zebra_crossing', 'crossroad', 'road', 'street', 'stairs', 'left t-junction', 't-junction'],
        conf_threshold=0.5,
        iou_threshold=0.5
    )
    
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    print("开始实时检测，按'q'键退出...")
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法获取帧")
            break
        
        # 执行检测
        detections = detector.detect(frame)
        
        # 可视化结果
        result_frame = detector.visualize(frame, detections)
        
        # 显示结果
        cv2.imshow('YOLOv8 实时检测', result_frame)
        
        # 按'q'键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()